# from Plate_Detection.MTCNN.Load_data.DataLoader import PlateDataset
from numpy.core.fromnumeric import _all_dispatcher
from torch.optim import optimizer
import os
import torch
from torch.utils.data import DataLoader
import sys
sys.path.append('..')
from Load_data.PlateData import *   # 载入数据集
from Core.model import *    # 载入模型
import cv2
import matplotlib.pyplot as plt
import time

model_store_path = os.path.join(os.path.dirname(os.getcwd()), 'Model')  # 倒时候要改

def compute_accuracy():
    pass

def train_pnet(model_store_path, batch_size, file_name, epoch, lr=0.01, use_cuda=True):

    if not os.path.exists(model_store_path):
        os.makedirs(model_store_path)

    # 损失函数
    lossfn = LossFn()

    # 网络
    net = PNet()
    net.train()

    if use_cuda:
        net.cuda()

    # 优化器
    optimizer = torch.optim.SGD(net.parameters(), lr=lr)

    # 载入数据集
    # image_12_47.txt or image_24_94.txt
    data = PlateDataset(file_name)

    # 用Dataloader载入数据
    PlateDataloader = DataLoader(data, batch_size=batch_size)   # 在Dataset内已经实现shuffle操作

    # 开始训练
    for cur_epoch in range(1, epoch+1): # 每一轮
        
        for image, gt_cls, gt_bbox in PlateDataloader:
            cur_time = time.time()
            count = 0
            # image：tensor
            if use_cuda:
                image = image.cuda()
                gt_cls = gt_cls.cuda()
                gt_bbox = gt_bbox.cuda()
            # 预测
            cls_pred, offset_pred = net(image)
            # loss
            cls_loss = lossfn.cls_loss(gt_cls, cls_pred)
            offset_loss = lossfn.box_loss(gt_cls, gt_bbox, offset_pred)
            all_loss = cls_loss * 1.0 + offset_loss * 0.5

            if count % 1000 == 0:
                end_time = time.time()
                used_time = end_time - cur_time
                # accuracy = compute_accuracy()
                # cls_acc = accuracy.cpu().numpy()
                cls_loss = cls_loss.cpu().numpy()
                offset_loss = offset_loss.cpu().numpy()
                all_loss = all_loss.cpu().numpy()
            count += 1
            print(f'epoch:{cur_epoch}, train:{count}, cls_loss:{cls_loss}, offset_loss:{offset_loss}, all_loss:{all_loss}, time:{used_time}')
            optimizer.zero_grad()
            all_loss.backward()
            optimizer.step()

        torch.save(net.state_dict(), os.path.join(model_store_path,"pnet_epoch_%d.pt" % cur_epoch))
        torch.save(net, os.path.join(model_store_path,"pnet_epoch_model_%d.pkl" % cur_epoch))
            

if __name__ == '__main__':
    train_pnet(model_store_path, batch_size=2, file_name='image_12_47.txt', epoch=5)
            

            


# PlateDataloader = DataLoader(PlateDataset('image_12_47.txt'), batch_size=2)
